# -*- coding:utf-8 -*-
import os
import sys
import time
import warnings
import logging
import numpy as np
import pandas as pd
from catboost import CatBoostClassifier

# 忽略 FutureWarning（仅限开发阶段）
warnings.filterwarnings("ignore", category=FutureWarning)

# 配置日志输出
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

# 获取当前脚本所在目录作为项目根目录
PROJECT_ROOT = os.path.dirname(os.path.abspath(__file__))


def validate_data_directory():
    """检查数据目录是否存在且包含必需的文件"""
    data_dir = os.path.join(PROJECT_ROOT, 'data', 'data_format1')
    required_files = [
        'test_format1.csv',
        'train_format1.csv',
        'user_info_format1.csv',
        'user_log_format1.csv'
    ]

    # 检查目录是否存在
    if not os.path.exists(data_dir):
        logger.error(f"数据目录不存在: {data_dir}")
        return False

    # 检查文件是否存在
    missing_files = []
    for file in required_files:
        file_path = os.path.join(data_dir, file)
        if not os.path.exists(file_path):
            missing_files.append(file_path)

    if missing_files:
        logger.error("以下数据文件缺失:")
        for file in missing_files:
            logger.error(f"  - {file}")
        return False

    logger.info(f"数据目录验证通过: {data_dir}")
    return True


def validate_model_directory():
    """检查模型目录是否存在且包含模型文件"""
    model_dir = os.path.join(PROJECT_ROOT, 'model')

    # 检查目录是否存在
    if not os.path.exists(model_dir):
        logger.error(f"模型目录不存在: {model_dir}")
        return None

    # 查找模型文件
    model_files = [f for f in os.listdir(model_dir) if f.endswith('.model')]

    if not model_files:
        logger.error(f"在模型目录中找不到任何 .model 文件: {model_dir}")
        return None

    # 按时间排序获取最新模型
    model_files.sort(reverse=True)
    latest_model = model_files[0]
    model_path = os.path.join(model_dir, latest_model)

    if not os.path.exists(model_path):
        logger.error(f"模型文件不存在: {model_path}")
        return None

    logger.info(f"找到模型文件: {model_path}")
    return model_path


def main():
    try:
        # 验证数据目录
        if not validate_data_directory():
            logger.error("数据目录验证失败，程序终止")
            return

        # 验证模型目录
        model_path = validate_model_directory()
        if not model_path:
            logger.error("模型文件验证失败，程序终止")
            return

        # 动态导入 dataset 模块
        import dataset as ds

        logger.info("加载测试数据集...")
        test_data, info_list = ds.load_test()
        logger.info(f"数据集加载完成，形状: {test_data.shape}")

        # 检查数据集是否为空
        if test_data.empty:
            logger.error("加载的测试数据集为空")
            return

        logger.info("填充缺失值...")
        test_data.fillna(0, inplace=True)  # 填充缺失值为0
        logger.info("填充完成")

        logger.info("加载模型...")
        model = CatBoostClassifier()
        model.load_model(model_path)
        logger.info(f"模型加载成功: {os.path.basename(model_path)}")

        logger.info("进行预测...")
        start_time = time.time()
        y_pred = model.predict_proba(test_data)
        ans = y_pred[:, 1].reshape(-1, 1)
        prediction_time = time.time() - start_time
        logger.info(f"预测完成，耗时: {prediction_time:.2f}秒")

        logger.info("生成答案文件...")
        # 确保info_list是二维数组
        if info_list.ndim == 1:
            info_list = info_list.reshape(-1, 1)

        # 合并结果
        answer = np.concatenate((info_list, ans), axis=1)
        answer_df = pd.DataFrame(answer, columns=["user_id", "merchant_id", "prob"])

        # 创建结果文件名
        ans_name = 'ans_' + time.strftime('%Y%m%d_%H%M%S', time.localtime())

        # 确保结果目录存在
        result_dir = os.path.join(PROJECT_ROOT, 'result')
        os.makedirs(result_dir, exist_ok=True)
        output_path = os.path.join(result_dir, ans_name + '.csv')

        logger.info(f"保存结果到: {output_path}")
        answer_df.to_csv(output_path, sep=',', header=True, index=False)
        logger.info("保存完成")
        logger.info("程序执行成功")

    except Exception as e:
        logger.error(f"发生错误: {e}", exc_info=True)
        logger.error("程序执行失败")


if __name__ == '__main__':
    # 添加项目根目录到系统路径
    sys.path.insert(0, PROJECT_ROOT)

    logger.info(f"项目根目录: {PROJECT_ROOT}")
    logger.info("启动天猫复购预测程序")
    main()
    logger.info("程序结束")